fix(ai-amazon-bedrock): emit finish part after metadata event to preserve token usage in streaming - #6160
Conversation
…erve token usage in streaming In the Bedrock ConverseStream API, events arrive in the order: messageStop → metadata. The metadata event carries all token usage data (inputTokens, outputTokens, totalTokens, cacheReadInputTokens, cacheWriteInputTokens) as well as trace information. Previously, the streaming handler emitted the finish part during the messageStop event, before the metadata event had populated the usage object. Because the downstream LanguageModel layer decodes stream parts through a schema (creating new instances), the finish part was captured with uninitialized usage values — causing all token counts (including cachedInputTokens) to be lost in streaming responses. The fix introduces a `tryEmitFinish` guard that defers the finish part until both messageStop (which provides the stop reason) and metadata (which provides usage/trace) have been received. This: - Ensures cachedInputTokens, inputTokens, outputTokens, totalTokens are all correctly populated in streaming responses - Ensures trace and cacheWriteInputTokens metadata are present - Handles event ordering defensively (emits on whichever arrives second) - Preserves correct behavior when errors interrupt the stream between messageStop and metadata (no misleading finish part with empty usage) The non-streaming (Converse) path was already correct since it reads directly from the decoded ConverseResponse. Adds comprehensive test suite for AmazonBedrockLanguageModel covering both streaming and non-streaming paths: token usage, cached tokens, explicit zero cache counts, trace metadata, tool calls, error scenarios, and missing event edge cases.
🦋 Changeset detectedLatest commit: 6995efe The changes in this PR will be included in the next version bump. This PR includes changesets to release 1 package
Not sure what this means? Click here to learn what changesets are. Click here if you're a maintainer who wants to add another changeset to this PR |
|
@jamesone - I think this PR needs a |
|
@loklaan - I'm not sure what changes were made but CI is failing. Looks like it might have to do with changes made to the lockfile. |
|
Woops missed that ping from Maxwell - @jamesone I think it was intended for you 🙇 gl |
|
Thanks for the PR. CI is currently failing; could you please update the branch and address the failing checks, or confirm whether this is still active? |
|
Thanks for the PR. The core Bedrock streaming finish/metadata fix appears to already be present on current main, and this branch is now conflicting with failing CI. I’m going to close this out; if there is remaining test coverage or behavior not covered by main, please open a fresh PR against current main. |
4 participants
Type
Description
In the Bedrock ConverseStream API, events arrive in the order:
messageStop→metadata. Themetadataevent carries all token usage data (inputTokens,outputTokens,totalTokens,cacheReadInputTokens,cacheWriteInputTokens) as well as trace information.Previously, the streaming handler emitted the finish part during the
messageStopevent, before themetadataevent had populated the usage object. Because the downstreamLanguageModellayer decodes stream parts through a schema (creating new instances), the finish part was captured with uninitialized usage values — causing all token counts (includingcachedInputTokens) to be lost in streaming responses.Fix
The fix introduces a
tryEmitFinishguard that defers the finish part until bothmessageStop(which provides the stop reason) andmetadata(which provides usage/trace) have been received. This:cachedInputTokens,inputTokens,outputTokens,totalTokensare all correctly populated in streaming responsestraceandcacheWriteInputTokensmetadata are presentmessageStopandmetadata(no misleading finish part with empty usage)The non-streaming (
Converse) path was already correct since it reads directly from the decodedConverseResponse.Tests
Adds comprehensive test suite for
AmazonBedrockLanguageModelcovering both streaming and non-streaming paths:cacheReadInputTokens/cacheWriteInputTokens)0(notundefined)messageStop(no misleading finish part)messageStopedge case (no finish part without a stop reason)Related